import os
import joblib
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.metrics import classification_report, confusion_matrix, roc_auc_score, roc_curve
from utils.data_load import dataload
from utils.feature_engineering_xgb import feature_engineering

def test_model():
    # 1️⃣ 加载数据
    X_train, X_test, y_train, y_test = dataload()

    # 2️⃣ 加载编码器和标准化器
    model = joblib.load("../model/xgboost_model.pkl")

    # 3️⃣ 应用特征工程（与训练一致）
    X_train_scaled, X_test_scaled, y_train_res, _, _ = feature_engineering(X_train, X_test, y_train)

    # 4️⃣ 模型预测
    y_pred = model.predict(X_test_scaled)
    y_proba = model.predict_proba(X_test_scaled)[:, 1]

    # 5️⃣ 输出结果
    print("\n===== 分类报告 =====")
    print(classification_report(y_test, y_pred))

    cm = confusion_matrix(y_test, y_pred)
    auc = roc_auc_score(y_test, y_proba)
    print("\n===== 混淆矩阵 =====")
    print(cm)
    print(f"\nROC AUC: {auc:.4f}")

    # 6️⃣ 可视化并保存
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', ax=axes[0])
    axes[0].set_title("Confusion Matrix")
    axes[0].set_xlabel("Predicted")
    axes[0].set_ylabel("True")

    fpr, tpr, _ = roc_curve(y_test, y_proba)
    axes[1].plot(fpr, tpr, color='darkorange', lw=2, label=f'AUC = {auc:.3f}')
    axes[1].plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    axes[1].set_title("ROC Curve")
    axes[1].set_xlabel("False Positive Rate")
    axes[1].set_ylabel("True Positive Rate")
    axes[1].legend(loc="lower right")

    plt.tight_layout()

    os.makedirs("../fig", exist_ok=True)
    fig_path = "../fig/xgboost_eval.png"
    plt.savefig(fig_path, dpi=300, bbox_inches='tight')
    print(f"✅ 图像已保存到: {fig_path}")

    plt.show()

if __name__ == "__main__":
    test_model()
